import numpy as np
import matplotlib.pyplot as plt
import os

def compute_dft_magnitude(signal, fs):
    """
    计算离散信号的傅里叶变换幅值
    :param signal: 输入信号序列 (1D np.array)
    :param fs: 信号的采样频率
    :return: 信号的正频率部分的幅值
    """
    # 1. 移除 DC 分量
    signal_zero_mean = signal - np.mean(signal)

    # 2. 计算 FFT
    signal_fre = np.fft.fft(signal_zero_mean)

    # 3. 仅保留正频率分量
    n = len(signal)  # 信号长度
    freqs = np.fft.fftfreq(n, d=1/fs)  # 频率轴，假设采样间隔为1/fs
    positive_indices = freqs >= 0  # 筛选正频率部分

    freqs_positive = freqs[positive_indices]
    signal_fre_positive = signal_fre[positive_indices]

    # 4. 计算幅值
    magnitude = np.abs(signal_fre_positive)

    return magnitude


def compute_y_over_x_magnitude(x_signal, y_signal, fs):
    """
    计算 Y/X 的傅里叶变换幅值
    :param x_signal: 输入 X 信号序列 (1D np.array)
    :param y_signal: 输入 Y 信号序列 (1D np.array)
    :param fs: 信号的采样频率
    :return: Y/X 的正频率部分的幅值
    """
    # 1. 移除 DC 分量
    x_signal_zero_mean = x_signal - np.mean(x_signal)
    y_signal_zero_mean = y_signal - np.mean(y_signal)

    # 2. 计算 X 和 Y 的 FFT
    x_signal_fre = np.fft.fft(x_signal_zero_mean)
    y_signal_fre = np.fft.fft(y_signal_zero_mean)

    # 3. 计算 Y/X
    threshold = 1e-6  # 阈值，避免除以过小的值
    y_over_x_fre = np.divide(y_signal_fre, x_signal_fre, out=np.zeros_like(y_signal_fre), where=np.abs(x_signal_fre) > threshold)

    # 4. 仅保留正频率分量
    n = len(x_signal)  # 信号长度
    freqs = np.fft.fftfreq(n, d=1/fs)  # 频率轴，假设采样间隔为1/fs
    positive_indices = freqs >= 0  # 筛选正频率部分

    freqs_positive = freqs[positive_indices]
    y_over_x_fre_positive = y_over_x_fre[positive_indices]

    # 5. 计算幅值
    magnitude = np.abs(y_over_x_fre_positive)

    return magnitude


def compute_and_plot_Y_over_X(a_x, a_y, b_x, b_y, fs, folder_a, folder_b, merged_folder_path, merged_folder, label_a, label_b, plot_title):
    """
    计算并绘制 Y/X 的幅值和相位图，并保存为图片。

    参数:
    - a_x, a_y: 信号 A 的 x 和 y 分量
    - b_x, b_y: 信号 B 的 x 和 y 分量
    - fs: 采样频率
    - folder_a, folder_b: 文件夹名称，用于绘图标签
    - merged_folder_path: 保存图片的文件夹路径
    - merged_folder: 合并文件夹的名称，用于生成文件名
    - label_a, label_b: 用于 A 和 B 的标签
    - plot_title: 图的标题
    """
    # 1. 移除 DC 分量
    a_x_zero_mean = a_x - np.mean(a_x)
    a_y_zero_mean = a_y - np.mean(a_y)
    b_x_zero_mean = b_x - np.mean(b_x)
    b_y_zero_mean = b_y - np.mean(b_y)

    # 2. 计算 FFT
    a_X_fre = np.fft.fft(a_x_zero_mean)
    a_Y_fre = np.fft.fft(a_y_zero_mean)
    b_X_fre = np.fft.fft(b_x_zero_mean)
    b_Y_fre = np.fft.fft(b_y_zero_mean)

    # 3. 计算 Y/X
    threshold = 1e-6  # 阈值，避免除以过小的值
    a_Y_over_X = np.divide(a_Y_fre, a_X_fre, out=np.zeros_like(a_Y_fre), where=np.abs(a_X_fre) > threshold)
    b_Y_over_X = np.divide(b_Y_fre, b_X_fre, out=np.zeros_like(b_Y_fre), where=np.abs(b_X_fre) > threshold)

    # 4. 仅保留正频率分量
    n = len(a_x)  # 信号长度
    freqs = np.fft.fftfreq(n, d=1/fs)  # 频率轴，假设采样间隔为1
    positive_indices = freqs >= 0  # 筛选正频率部分

    freqs_positive = freqs[positive_indices]
    a_Y_over_X_positive = a_Y_over_X[positive_indices]
    b_Y_over_X_positive = b_Y_over_X[positive_indices]

    # 5. 绘制结果
    # 绘制 Y/X 的幅值
    plt.figure(figsize=(12, 6))  # 调整图像大小

    # 绘制幅值
    plt.plot(freqs_positive, np.abs(a_Y_over_X_positive), label=f'{label_a} Magnitude ({folder_a})', color='blue', linewidth=1)
    plt.scatter(freqs_positive, np.abs(a_Y_over_X_positive), color='blue', s=15, zorder=3)  # 增大散点标记大小
    plt.plot(freqs_positive, np.abs(b_Y_over_X_positive), label=f'{label_b} Magnitude ({folder_b})', color='orange', linewidth=1)
    plt.scatter(freqs_positive, np.abs(b_Y_over_X_positive), color='orange', s=15, zorder=3)  # 增大散点标记大小

    # 添加标题和标签
    plt.title(f'Magnitude of {plot_title}', fontsize=14)  # 图的标题
    plt.xlabel("Frequency (Hz)", fontsize=12)
    plt.ylabel("Magnitude", fontsize=12)
    plt.grid(alpha=0.5)  # 调整网格透明度
    plt.legend(fontsize=12)

    # 保存幅值图片
    plt.savefig(os.path.join(merged_folder_path, f"{merged_folder}_{plot_title}_magnitude.png"), dpi=300)
    plt.close()
    print(f"Image saved as {merged_folder}_{plot_title}_magnitude.png")
